from __future__ import absolute_import, unicode_literals
from celery import shared_task

from django.db.models import Q

from hte.error.handle import abort_on_error
from hte.error.models import HTEError

from .handle import handle_ml
from .models.task import Task
from .pic_recog.handle import pic_training, test_images
from .models.task import VisualRecognitionTask
from .models.vr_model import VisualRecognitionModel
from .models.asset import ClassAsset


@shared_task
def async_ml(task_id):
    task = Task.objects.get(pk=task_id)
    result = handle_ml(task)
    return result


@shared_task
def visual_recognition_task(vrtask_id, pic_list, classes_list, test_dir):
    '''

    :param vrtask_id: 用户创建的图像识别任务的ID
    :param pic_list: 训练集的图片路径，list
    :param classes_list: 对应pic_list的分类列表，list
    :param test_dir: 测试集的图片路径列表，list
    :return:
    '''
    train_result, modle_dir, name_to_label = pic_training(pic_list=pic_list, classes=classes_list)
    # test_result = test_images(modle_dir['train_dir'], test_dir, name_to_label)
    #  这里因为在api_v2中调用api的时候已经检验过ID错误，任务不存在的情况，因此这里不用再做检查
    vr_task = VisualRecognitionTask.objects.get(id=vrtask_id)
    vr_task.train_result = train_result
    vr_task.model_dir = modle_dir
    vr_task.name_to_label = name_to_label
    #  在这里直接就开始测试
    test_result = test_images(modle_dir['train_dir'], test_dir, name_to_label)
    vr_task.test_result = test_result
    # 这里添加任务执行情况的判定，并写入到数据库
    if test_result is not None:
        vr_task.status = 2
    else:
        vr_task.status = 3
    try:
        vr_task.save()
    except Exception:
        raise
    # 在这里添加模型
    author = vr_task.author
    q = Q(vrtask_id=vrtask_id) & ~Q(title='test')
    class_assets = ClassAsset.objects.filter(q)
    classes_images = list()
    # images_of_classes = list()
    for class_asset in class_assets:
        temp = dict()
        temp['class_title'] = class_asset.title
        temp['num_of_pictures'] = class_asset.num_of_pictures
        classes_images.append(temp)

    vr_model = VisualRecognitionModel(author=author, classes_images=classes_images,
                                      vrtask_id=vrtask_id, path=modle_dir, result=train_result)
    try:
        vr_model.save()
    except Exception:
        raise
    return test_result


@shared_task
def add(x, y):
    return x + y


@shared_task
def mul(x, y):
    return x * y
